import os
import re
from datetime import datetime
from src.core.database import DatabaseConnector
from src.models.owner import Owner, OwnerRelation
from src.io.file_exporter import FileExporter
from src.io.csv_importer import CsvImporter
from src.core.dbhelper import DBHelper

class OwnerCommand:
    """单位命令处理类"""
    
    def __init__(self, db_connector=None, subparsers=None):
        """初始化单位命令处理器"""
        self.db_connector = db_connector or DatabaseConnector()
        self.file_exporter = FileExporter()
        self.csv_importer = CsvImporter(self.db_connector)
        self.db_helper = DBHelper(self.db_connector)
        self.arg_parser = self.setup_parser(subparsers)

    def setup_parser(self, subparsers):
        """设置单位管理命令行参数解析器"""
        # owner命令组
        owner_parser = subparsers.add_parser('owner', help='单位管理相关命令', description='单位管理相关的所有操作')
        owner_subparsers = owner_parser.add_subparsers(dest='owner_command', help='单位管理子命令')

        # owner help命令组
        owner_subparsers.add_parser('help', help='显示单位管理命令帮助')

        # owner add命令组
        owner_add_parser = owner_subparsers.add_parser('add', help='添加单位')
        owner_add_parser.add_argument('name', help='单位名称')
        owner_add_parser.add_argument('--parent', help='父单位ID或名称')
        owner_add_parser.add_argument('--parent-relation', help='与父单位的关系类型')

        # owner get命令组
        owner_get_parser = owner_subparsers.add_parser('get', help='获取单位信息')
        owner_get_parser.add_argument('identifier', help='单位ID或名称')
        owner_get_parser.add_argument('relation_type', nargs='?', choices=['ip', 'domain', 'parent', 'child'], help='关系类型')
        owner_get_parser.add_argument('-o', '--output', nargs='?', const=True, help='指定关系类型时，导出文件路径，不指定则自动生成')
        owner_get_parser.add_argument('-e', '--extension', choices=['json', 'csv'], default='csv', help='导出文件格式，默认csv')

        # owner set命令组
        owner_set_parser = owner_subparsers.add_parser('set', help='修改单位名称')
        owner_set_parser.add_argument('identifier', help='单位ID或名称')
        owner_set_parser.add_argument('--name', help='新的单位名称')
        owner_set_parser.add_argument('--parent', help='新的父单位ID或名称')
        owner_set_parser.add_argument('--parent-relation', help='与新父单位的关系类型')

        # owner getall命令组
        owner_getall_parser = owner_subparsers.add_parser('getall', help='获取所有单位列表')
        # 设置--export为可选参数，不提供参数值时自动生成文件名
        owner_getall_parser.add_argument('-o', '--output', nargs='?', const=True, help='导出文件路径，不指定则自动生成')
        owner_getall_parser.add_argument('-e', '--extension', choices=['json', 'csv'], default='csv', help='导出文件格式，默认csv')

        # owner search命令组
        owner_search_parser = owner_subparsers.add_parser('search', help='搜索单位')
        owner_search_parser.add_argument('keyword', help='搜索关键词')
        owner_search_parser.add_argument('-o', '--output', nargs='?', const=True, help='导出文件路径，不指定则自动生成')
        owner_search_parser.add_argument('-e', '--extension', choices=['json', 'csv'], default='csv', help='导出文件格式，默认csv')

        # owner delete命令组
        owner_delete_parser = owner_subparsers.add_parser('delete', help='删除单位')
        owner_delete_parser.add_argument('identifier', help='单位ID或名称')

        # owner import命令组
        owner_import_parser = owner_subparsers.add_parser('import', help='导入单位数据')
        owner_import_parser.add_argument('file_path', help='CSV文件路径')
        owner_import_parser.add_argument('--name-column', help='单位名称列的索引或名称，默认为表头为"owner"的列，或第一列', default='owner')

        return owner_parser

    def handle_command(self, args):
        """处理单位命令"""
        command = args.owner_command
        if command == 'help':
            self.show_owner_help()
        elif command == 'add':
            if not hasattr(args, 'name'):
                print("添加单位参数不足！格式: owner add <单位名称>")
                return
            parent = args.parent if hasattr(args, 'parent') else None
            parent_relation = args.parent_relation
            self.add_owner(args.name, parent, parent_relation)
        elif command == 'get':
            if not args.identifier:
                print("获取单位参数不足！格式: owner get <单位ID/名称>")
                return
            relation_type = args.relation_type
            export_path = args.output if hasattr(args, 'export') else None
            extension = args.extension if hasattr(args, 'extension') else 'csv'
            self.get_owner(args.identifier, relation_type, export_path=export_path, extension=extension)
        elif command == 'set':
            if not hasattr(args, 'identifier'):
                print("修改单位参数不足！格式: owner set <单位ID/名称> [参数]")
                return
            name = args.name
            parent = args.parent
            parent_type = args.parent_relation
            if name:
                self.set_owner_name(args.identifier, name)
            if parent is not None:
                self.set_owner_parent(args.identifier, parent, parent_type)
        elif command == 'getall':
            export_path = args.output if hasattr(args, 'export') else None
            extension = args.extension if hasattr(args, 'extension') else 'csv'
            self.get_all_owners(export_path=export_path, extension=extension)
        elif command == 'search':
            if not hasattr(args, 'keyword'):
                print("搜索单位参数不足！格式: owner search <关键词>")
                return
            export_path = args.output if hasattr(args, 'export') else None
            extension = args.extension if hasattr(args, 'extension') else 'csv'
            self.search_owners(args.keyword, export_path, extension)
        elif command == 'delete':
            if not hasattr(args, 'identifier'):
                print("删除单位参数不足！格式: owner delete <单位ID/名称>")
                return
            self.delete_owner(args.identifier)
        elif command == 'import':
            if not hasattr(args, 'file_path'):
                print("导入单位参数不足！格式: owner import <file_path>")
                return
            name_column = args.name_column if hasattr(args, 'name_column') else 'owner'
            self.import_owners(args.file_path, name_column)
        else:
            print(f"未知的单位命令: {command}")
            self.show_owner_help()
    
    def show_owner_help(self):
        """显示单位管理命令帮助信息"""
        self.arg_parser.print_help()
    
    def add_owner(self, owner_name, parent_identifier=None, parent_relation=None):
        """添加新单位"""
        owner = self.db_helper.get_or_add_owner(owner_name, parent_identifier, parent_relation)
        if '_is_new' in owner and owner['_is_new']:
            print(f"成功添加新单位: ID={owner['owner_id']}")
            print(f"  名称: {owner['owner_name']}")
        else:
            print(f"单位 {owner_name} 已存在，ID={owner['owner_id']}")
        
    
    def get_owner(self, identifier, relation_type=None, export_path=None, extension=None):
        """获取指定单位信息"""
        owner = self._get_owner(identifier)
        if not owner:
            return
            
        print("单位信息:")
        print(f"  ID: {owner['owner_id']}")
        print(f"  名称: {owner['owner_name']}")
        print(f"  创建时间: {owner['create_time']}")
        print(f"  更新时间: {owner['update_time']}")

        print("")
        if relation_type == 'parent':
            self.get_owner_parents(owner, export_path=export_path, extension=extension)
        elif relation_type == 'child':
            self.get_owner_children(owner, export_path=export_path, extension=extension)
        elif relation_type == 'ip':
            self.get_owner_ips(owner, export_path=export_path, extension=extension)
        elif relation_type == 'domain':
            self.get_owner_domains(owner, export_path=export_path, extension=extension)
    
    def get_all_owners(self, export_path=None, extension=None):
        """获取所有单位列表"""
        # 查询所有单位
        sql = "SELECT * FROM owners ORDER BY owner_id"
        cursor = self.db_connector.execute_query(sql)
        
        if not cursor:
            print("暂无单位信息！")
            return
        
        # 获取结果列表
        if hasattr(cursor, 'fetchall'):
            owners_result = cursor.fetchall()
        else:
            owners_result = cursor
            
        if not owners_result:
            print("暂无单位信息！")
            return
        
        # 准备导出数据
        export_data = []
        for owner in owners_result:
            export_item = owner.copy() if isinstance(owner, dict) else {
                'owner_id': owner[0],
                'owner_name': owner[1],
                'create_time': owner[2],
                'update_time': owner[3]
            }
            # 转换datetime对象为字符串
            if 'create_time' in export_item and isinstance(export_item['create_time'], datetime):
                export_item['create_time'] = export_item['create_time'].strftime('%Y-%m-%d %H:%M:%S')
            if 'update_time' in export_item and isinstance(export_item['update_time'], datetime):
                export_item['update_time'] = export_item['update_time'].strftime('%Y-%m-%d %H:%M:%S')
            export_data.append(export_item)
        
        # 导出数据
        if export_path is not None:
            actual_path = None if export_path is True else export_path
            self.file_exporter.export_data(export_data, actual_path, 'all_owners', extension)
            return
        
        # 显示单位列表
        print(f"共有 {len(owners_result)} 个单位:")
        print("{:<10} {:<30} {:<20} {:<20}".format(
            "ID", "单位名称", "创建时间", "更新时间"
        ))
        print("=" * 80)
        
        for row in owners_result:
            if isinstance(row, (list, tuple)):
                owner_id = row[0]
                owner_name = row[1]
                create_time = row[2]
                update_time = row[3]
            elif isinstance(row, dict):
                owner_id = row['owner_id']
                owner_name = row['owner_name']
                create_time = row['create_time']
                update_time = row['update_time']
            
            # 转换datetime对象为字符串
            if isinstance(create_time, datetime):
                create_time = create_time.strftime('%Y-%m-%d %H:%M:%S')
            if isinstance(update_time, datetime):
                update_time = update_time.strftime('%Y-%m-%d %H:%M:%S')
            
            print("{:<10} {:<30} {:<20} {:<20}".format(
                owner_id, owner_name, create_time, update_time
            ))
        print(f"共有 {len(owners_result)} 个单位")
    
    def search_owners(self, keyword, export_path=None, extension=None):
        """搜索单位"""
        # 构建搜索条件
        sql = "SELECT * FROM owners WHERE owner_name LIKE %s ORDER BY owner_id"
        params = [f"%{keyword}%"]
        
        cursor = self.db_connector.execute_query(sql, params)
        
        if not cursor:
            print(f"未找到包含 '{keyword}' 的单位！")
            return
        
        # 获取结果列表
        if hasattr(cursor, 'fetchall'):
            owners_result = cursor.fetchall()
        else:
            owners_result = cursor
            
        if not owners_result:
            print(f"未找到包含 '{keyword}' 的单位！")
            return
        
        # 准备导出数据
        export_data = []
        for owner in owners_result:
            export_item = owner.copy() if isinstance(owner, dict) else {
                'owner_id': owner[0],
                'owner_name': owner[1],
                'create_time': owner[2],
                'update_time': owner[3]
            }
            # 转换datetime对象为字符串
            if 'create_time' in export_item and isinstance(export_item['create_time'], datetime):
                export_item['create_time'] = export_item['create_time'].strftime('%Y-%m-%d %H:%M:%S')
            if 'update_time' in export_item and isinstance(export_item['update_time'], datetime):
                export_item['update_time'] = export_item['update_time'].strftime('%Y-%m-%d %H:%M:%S')
            export_data.append(export_item)
        
        # 导出数据
        if export_path is not None:
            actual_path = None if export_path is True else export_path
            self.file_exporter.export_data(export_data, actual_path, 'owner_search', extension)
            return
        
        # 显示搜索结果
        print(f"找到 {len(owners_result)} 个匹配的单位:")
        print("{:<10} {:<30} {:<20} {:<20}".format(
            "ID", "单位名称", "创建时间", "更新时间"
        ))
        print("=" * 80)
        
        for row in owners_result:
            if isinstance(row, (list, tuple)):
                owner_id = row[0]
                owner_name = row[1]
                create_time = row[2]
                update_time = row[3]
            elif isinstance(row, dict):
                owner_id = row['owner_id']
                owner_name = row['owner_name']
                create_time = row['create_time']
                update_time = row['update_time']
            
            # 转换datetime对象为字符串
            if isinstance(create_time, datetime):
                create_time = create_time.strftime('%Y-%m-%d %H:%M:%S')
            if isinstance(update_time, datetime):
                update_time = update_time.strftime('%Y-%m-%d %H:%M:%S')
            
            print("{:<10} {:<30} {:<20} {:<20}".format(
                owner_id, owner_name, create_time, update_time
            ))
    
    def delete_owner(self, identifier):
        """删除单位"""
        # 查找单位
        try:
            owner_id = int(identifier)
            sql = "SELECT owner_id, owner_name FROM owners WHERE owner_id = %s"
            params = [owner_id]
        except ValueError:
            sql = "SELECT owner_id, owner_name FROM owners WHERE owner_name = %s"
            params = [identifier]
        
        result = self.db_connector.fetchone(sql, params)
        if not result:
            print(f"未找到单位: {identifier}")
            return
        
        # 获取单位ID和名称
        owner_id = result['owner_id']
        owner_name = result['owner_name']
        
        # 检查是否与其他单位有父子关系
        sql = "SELECT COUNT(*) as count FROM owner_relations WHERE parent_id = %s or child_id = %s"
        parent_count = self.db_connector.fetchvalue(sql, (owner_id, owner_id))
        if parent_count > 0:
            print(f"错误：单位 '{owner_name}' 有 {parent_count} 个子单位！")
            return

        # 检查是否域名或IP关联了此单位
        sql = "SELECT COUNT(*) as count FROM domains WHERE owner_id = %s"
        domain_count = self.db_connector.fetchvalue(sql, (owner_id,))
        if domain_count > 0:
            print(f"错误：单位 '{owner_name}' 关联了 {domain_count} 个域名！")
            return
        sql = "SELECT COUNT(*) as count FROM ips WHERE owner_id = %s"
        ip_count = self.db_connector.fetchvalue(sql, (owner_id,))
        if ip_count > 0:
            print(f"错误：单位 '{owner_name}' 关联了 {ip_count} 个IP地址！")
            return

        # 检查是否有资产关联到此单位
        sql = "SELECT COUNT(*) as count FROM assets WHERE owner_id = %s"
        asset_count = self.db_connector.fetchvalue(sql, (owner_id,))
        if asset_count > 0:
            print(f"错误：单位 '{owner_name}' 关联了 {asset_count} 个资产！")
            return
        
        # 删除单位
        sql = "DELETE FROM owners WHERE owner_id = %s"
        self.db_connector.execute_query(sql, (owner_id,))
        print(f"成功删除单位: ID={owner_id}, 名称='{owner_name}'")
    
    def get_owner_ips(self, owner, export_path=None, extension=None):
        """获取单位关联的IP地址"""
        # 查询关联的IP地址
        owner_id = owner['owner_id']
        owner_name = owner['owner_name']
        sql = "select * from ips where owner_id = %s"
        result = self.db_connector.fetchall(sql, (owner_id,))
        
        if len(result) == 0:
            print(f"单位ID {owner_id} 没有关联的IP地址！")
            return
        
        # 导出数据
        if export_path is not None:
            export_data = []
            for row in result:
                export_data.append({
                    'id': row['ip_id'],
                    'ip': row['ip'],
                    'owner_id': owner_id,
                    'owner_name': owner_name
                })
            actual_path = None if export_path is True else export_path
            self.file_exporter.export_data(export_data, actual_path, f'owner_{owner_id}_ips', extension)
            return
        
        # 显示IP列表
        print(f"单位ID {owner_id} 关联的IP地址 ({len(result)} 个):")
        print("{:<15} {:<10}".format("ID", "IP地址"))
        print("=" * 25)
        
        for row in result:
            print("{:<15} {:<10}".format(row['ip_id'], row['ip']))
    
    def get_owner_domains(self, owner, export_path=None, extension=None):
        """获取单位关联的域名"""
        # 查询关联的域名
        owner_id = owner['owner_id']
        owner_name = owner['owner_name']
        sql = "select * from domains where owner_id = %s"
        result = self.db_connector.fetchall(sql, (owner_id,))

        if len(result) == 0:
            print(f"单位ID {owner_id} 没有关联的域名！")
            return

        # 导出数据
        if export_path is not None:
            export_data = []
            for row in result:
                export_data.append({
                    'id': row['domain_id'],
                    'domain': row['domain'],
                    'owner_id': owner_id,
                    'owner_name': owner_name
                })
            actual_path = None if export_path is True else export_path
            self.file_exporter.export_data(export_data, actual_path, f'owner_{owner_id}_domains', extension)
            return

        # 显示域名列表
        print(f"单位ID {owner_id} 关联的域名 ({len(result)} 个):")
        print("{:<15} {:<10}".format("ID", "域名"))
        print("=" * 25)

        for row in result:
            print("{:<15} {:<10}".format(row['domain_id'], row['domain']))
    
    def import_owners(self, file_path, name_column=None):
        """从CSV文件导入单位数据"""
        # 验证文件是否存在
        if not os.path.exists(file_path):
            print(f"错误: 文件 {file_path} 不存在！")
            return
        
        # 检查文件扩展名
        file_ext = os.path.splitext(file_path)[1].lower()
        if file_ext != '.csv':
            print(f"错误: 只支持CSV文件格式，当前文件格式为 '{file_ext}'")
            return
        
        try:
            success_count, failed_count = self.csv_importer.import_owners_from_csv(file_path, name_column)
            print(f"单位导入完成: 成功 {success_count} 个, 失败 {failed_count} 个")
        except Exception as e:
            print(f"导入单位时发生错误: {str(e)}")
    
    def _get_owner(self, identifier):
        """获取单位"""
        return self.db_helper.get_owner(identifier)
            
    def set_owner_name(self, identifier, new_name):
        """修改单位名称"""
        # 获取单位ID
        owner = self._get_owner(identifier)
        if not owner:
            print(f"错误：未找到单位 '{identifier}'！")
            return
        
        owner_id = owner['owner_id']
        # 检查新名称是否已存在
        sql = "SELECT owner_id FROM owners WHERE owner_name = %s AND owner_id != %s"
        result = self.db_connector.fetchone(sql, (new_name, owner_id))
        if result:
            print(f"错误：名称 '{new_name}' 已被其他单位使用！")
            return

        if owner['owner_name'] == new_name:
            print(f"提示：单位 '{identifier}' 的名称未改变！")
            return
        
        # 更新单位名称
        current_time = datetime.now()
        sql = "UPDATE owners SET owner_name = %s, update_time = %s WHERE owner_id = %s"
        
        try:
            self.db_connector.execute_query(sql, (new_name, current_time, owner_id))
            print(f"成功修改单位ID {owner_id} 的名称 '{owner['owner_name']}' -> '{new_name}'")
        except Exception as e:
            print(f"修改单位名称失败: {e}")
    
    def get_owner_parents(self, owner, export_path=None, extension=None):
        """获取单位的父单位列表"""
        owner_id = owner['owner_id']
        owner_name = owner['owner_name']
        sql = "SELECT *, o.owner_name parent_name from owner_relations orl LEFT JOIN owners o ON orl.parent_id = o.owner_id WHERE child_id = %s"
        parents_result = self.db_connector.fetchall(sql, (owner_id,))
        if len(parents_result) == 0:
            print(f"单位ID {owner_id} 没有父单位！")
            return

        if export_path is not None:
            export_data = []
            for row in parents_result:
                export_data.append({
                    'owner_id': owner_id,
                    'owner_name': owner_name,
                    'parent_id': row['parent_id'],
                    'parent_name': row['parent_name'],
                    'relation': row['parent_type']
                })
            actual_path = None if export_path is True else export_path
            self.file_exporter.export_data(export_data, actual_path, f'owner_{owner_id}_parents', extension)
            return

        print(f"单位ID {owner_id} 的父单位列表 ({len(parents_result)} 个):")
        print("{:<10} {:<30} {:<10}".format("ID", "父单位名称", "父对子关系"))
        print("=" * 60)

        for row in parents_result:
            relation = row['parent_type'] if row['parent_type'] else ''
            print("{:<10} {:<30} {:<10}".format(row['owner_id'], row['parent_name'], relation))
    
    def get_owner_children(self, owner, export_path=None, extension=None):
        """获取单位的子单位列表"""
        owner_id = owner['owner_id']
        owner_name = owner['owner_name']
        sql = "SELECT *, o.owner_name child_name from owner_relations orl LEFT JOIN owners o ON orl.child_id = o.owner_id WHERE parent_id = %s"
        children_result = self.db_connector.fetchall(sql, (owner_id,))
        if len(children_result) == 0:
            print(f"单位ID {owner_id} 没有子单位！")
            return
        
        if export_path is not None:
            export_data = []
            for row in children_result:
                export_data.append({
                    'owner_id': owner_id,
                    'owner_name': owner_name,
                    'child_id': row['child_id'],
                    'child_name': row['child_name'],
                    'relation': row['parent_type']
                })
            actual_path = None if export_path is True else export_path
            self.file_exporter.export_data(export_data, actual_path, f'owner_{owner_id}_children', extension)
            return
        
        print(f"单位ID {owner_id} 的子单位列表 ({len(children_result)} 个):")
        print("{:<10} {:<30} {:<10}".format("ID", "子单位名称", "父对子关系"))
        print("=" * 60)

        for row in children_result:
            relation = row['parent_type'] if row['parent_type'] else ''
            print("{:<10} {:<30} {:<10}".format(row['owner_id'], row['child_name'], relation))
    
    def _has_circular_reference(self, parent_id, child_id):
        """检查是否存在循环引用"""
        # 简单的循环引用检查，通过递归查找子单位的父单位链
        visited = set()
        current_id = parent_id
        
        while current_id:
            if current_id == child_id:
                return True  # 发现循环引用
            if current_id in visited:
                return False  # 已经访问过，避免无限循环
                
            visited.add(current_id)
            
            # 查找当前单位的父单位
            sql = "SELECT parent_id FROM owner_relations WHERE child_id = %s LIMIT 1"
            cursor = self.db_connector.execute_query(sql, (current_id,))
            
            next_parent = None
            if hasattr(cursor, 'fetchone'):
                result = cursor.fetchone()
                if result:
                    next_parent = result[0] if isinstance(result, (list, tuple)) else result['parent_id']
            elif cursor and len(cursor) > 0:
                next_parent = cursor[0][0] if isinstance(cursor[0], (list, tuple)) else cursor[0]['parent_id']
                
            current_id = next_parent
            
        return False
    
    def _is_valid_ip(self, ip):
        """验证IP地址格式是否有效"""
        # 简单的IP地址格式验证
        pattern = r"^((25[0-5]|2[0-4]\d|[01]?\d\d?)\.){3}(25[0-5]|2[0-4]\d|[01]?\d\d?)$"
        return re.match(pattern, ip) is not None
        
    def set_owner_parent(self, owner_identifier, parent_identifier, parent_type='总公司'):
        """设置单位的父单位（会先删除已有的父单位关系）"""
        # 获取单位
        owner = self._get_owner(owner_identifier)
        if not owner:
            print(f"单位 {owner_identifier} 不存在！")
            return
        
        owner_id = owner['owner_id']
        owner_name = owner['owner_name']
        
        # 获取父单位
        parent = self._get_owner(parent_identifier)
        if not parent:
            print(f"父单位 {parent_identifier} 不存在！")
            return

        parent_id = parent['owner_id']
        parent_name = parent['owner_name']

        # 检查父单位是否为自己
        if owner_id == parent_id:
            print("错误：不能将自己设置为自己的父单位！")
            return
        
        # 检查是否存在循环引用
        if self._has_circular_reference(parent_id, owner_id):
            print("错误：添加此父单位会导致循环引用！")
            return
        
        # 删除现有的所有父单位关系
        sql_delete = "DELETE FROM owner_relations WHERE child_id = %s"
        try:
            self.db_connector.execute_query(sql_delete, (owner_id,))
        except Exception as e:
            print(f"删除现有父单位关系失败: {e}")
            return
        
        # 添加新的父子关系
        current_time = datetime.now()
        sql_insert = "INSERT INTO owner_relations (parent_id, child_id, create_time, update_time, parent_type, child_type) VALUES (%s, %s, %s, %s, %s, %s)"
        
        try:
            self.db_connector.execute_query(sql_insert, (parent_id, owner_id, current_time, current_time, parent_type, ''))
            print(f"成功设置父单位关系：单位ID {owner_id}({owner_name}) → 父单位ID {parent_id}({parent_name})")
        except Exception as e:
            print(f"设置父单位失败: {e}")